/**
 * Copyright (c) 2018-2023, NXOS Development Team
 * SPDX-License-Identifier: Apache-2.0
 *
 * Contains: Implementation of virtio-blk
 *
 * Change Logs:
 * Date           Author            Notes
 * 2023-07-03     Shiroko           Init
 */
#include "nxos.h"

#ifdef CONFIG_NX_DRIVER_VIRTIO_BLK

#undef NX_LOG_MOD_NAME
#define NX_LOG_MOD_NAME "VIRTIO-BLK"

#include "base/log.h"

#define DRV_NAME "virtio blk"
#define DEV_NAME "virtio_blk"

#include "../virtio.h"
#include "virtio_blk.h"
#include "base/block.h"


// Caller must ensure buffer is phy addr
NX_PRIVATE NX_Error
VirtioBLKReadWrite(NX_Device *device, NX_U64 sector, const char *buffer, NX_Size count, NX_Bool write)
{
    NX_U32      idx[3];
    VirtioQueue *queue = VirtioQueueGet(device, 0);
    // Allocate descriptor
    while (VirtioQueueDescAllocMany(queue, 3, idx) != NX_EOK);

    // Split the request form into three section: type+sector, data, status
    // Part 1
    VirtioBLKRequestPart1 part1 = {
            .type = write ? VIRTIO_BLK_TYPE_OUT : VIRTIO_BLK_TYPE_IN,
            .reserved = 0,
            .sector = sector
    };
    VirtioQueueDescFill(queue, idx[0], (NX_U64) ((NX_UArch) NX_Virt2Phy(&part1)), sizeof(part1),
                        VIRTIO_QUEUE_DESC_FLAG_NEXT, idx[1]);
    // Part 2
    VirtioQueueDescFill(queue, idx[1], (NX_U64) ((NX_UArch) buffer), count * VIRTIO_BLK_SECTOR_SIZE,
                        VIRTIO_QUEUE_DESC_FLAG_NEXT | (write ? 0 : VIRTIO_QUEUE_DESC_FLAG_WRITE), idx[2]);
    // Part 3
    NX_U8 status = 0;
    VirtioQueueDescFill(queue, idx[2], (NX_U64) ((NX_UArch) NX_Virt2Phy(&status)), 1, VIRTIO_QUEUE_DESC_FLAG_WRITE, 0);

    // Put the descriptor into available ring
    VirtioQueuePutIntoAvailRing(queue, idx[0]);
    // Notify device about queue
    VirtioSendNotification(device, 0, idx[0]);
    // Wait for complete
    VirtioWaitNotification(device, 0, idx[0]);
    // Free resources
    VirtioQueueDescFreeMany(queue, idx[0]);

    return NX_EOK;
}

NX_PRIVATE NX_Error VirtioBLKFlush(NX_Device *device)
{
    NX_U32      idx[2];
    VirtioQueue *queue = VirtioQueueGet(device, 0);
    // Allocate descriptor
    while (VirtioQueueDescAllocMany(queue, 2, idx) != NX_EOK);
    // Split the request form into three section: type+sector, data, status
    // Part 1
    VirtioBLKRequestPart1 part1 = {
            .type = VIRTIO_BLK_TYPE_FLUSH,
            .reserved = 0,
            .sector = 0
    };
    VirtioQueueDescFill(queue, idx[0], (NX_U64) ((NX_UArch) NX_Virt2Phy(&part1)), sizeof(part1),
                        VIRTIO_QUEUE_DESC_FLAG_NEXT, idx[1]);
    // Part 2
    NX_U8 status = 0;
    VirtioQueueDescFill(queue, idx[1], (NX_U64) ((NX_UArch) NX_Virt2Phy(&status)), 1,
                        VIRTIO_QUEUE_DESC_FLAG_WRITE, 0);

    // Put the descriptor into available ring
    VirtioQueuePutIntoAvailRing(queue, idx[0]);
    // Notify device about queue
    VirtioSendNotification(device, 0, idx[0]);
    // Wait for complete
    VirtioWaitNotification(device, 0, idx[0]);
    // Free resources
    VirtioQueueDescFreeMany(queue, idx[0]);

    return NX_EOK;
}

/* Driver Operations */
NX_PRIVATE NX_Error open(struct NX_Device *device, NX_U32 flags)
{
    if (VirtioGetExtension(device)->devExt != NX_NULL) return NX_EBUSY;
    NX_U32 features                    = 0;
    features &= ~(VIRTIO_BLK_FEATURE_RO | VIRTIO_BLK_FEATURE_MQ | VIRTIO_BLK_FEATURE_CONFIG_WCE);
    features &= ~(VIRTIO_BLK_FEATURE_SCSI);
    features &= VIRTIO_BLK_FEATURE_BLK_SIZE;

    if (flags & VIRTIO_BLK_FLAGS_READ_ONLY)
        features &= VIRTIO_BLK_FEATURE_RO;

    // Negotiate features & status
    VirtioDeviceInitialize(device, &features, 1);

    // Device-specified initialization
    VirtioQueueAlloc(device, 1);
    VirtioQueueInit(device, 0, CONFIG_NX_DRIVER_VIRTIO_SPLIT_QUEUE_SIZE);

    // Driver OK
    VirtioDeviceDriverOK(device);

    // Debug output device info
    VirtioBLKDeviceConfig *deviceConfig = NX_MemAlloc(sizeof(VirtioBLKDeviceConfig));
    if (deviceConfig == NX_NULL)
    {
        VirtioQueueDestroy(device, 0);
        VirtioQueueFree(device);
        return NX_ENOMEM;
    }
    VirtioReadDeviceConfig(device, 0, deviceConfig, sizeof(VirtioBLKDeviceConfig));
    NX_LOG_D("%s opened (Bus: %s), capacity: 0x%X bytes.", device->name,
             VirtioGetExtension(device)->busType == VIRTIO_BUS_TYPE_PCI ? "PCI" : "MMIO",
             deviceConfig->capacity * VIRTIO_BLK_SECTOR_SIZE);

    VirtioGetExtension(device)->flags  = flags;
    VirtioGetExtension(device)->devExt = deviceConfig;

    return NX_EOK;
}

NX_PRIVATE NX_Error close(struct NX_Device *device)
{
    // Reset the device
    VirtioDeviceReset(device);
    // Cleanup VirtioQueue
    for (int i                         = 0; i < VirtioGetExtension(device)->queue_num; i++)
        VirtioQueueDestroy(device, i);
    VirtioQueueFree(device);
    NX_MemFree(VirtioGetExtension(device)->devExt);
    VirtioGetExtension(device)->devExt = NX_NULL;
    return NX_EOK;
}

NX_PRIVATE NX_Error read(struct NX_Device *device, void *buf, NX_Offset off, NX_Size len, NX_Size *outLen)
{
    NX_ASSERT(buf != NX_NULL);
    NX_ASSERT(len != 0);
    buf = VirtioVirt2Phy(buf);
    if (off % VIRTIO_BLK_SECTOR_SIZE == 0 && len % VIRTIO_BLK_SECTOR_SIZE == 0)
    {
        // Perfect aligned, directly pass buffer to device without internal copy
        if (off + len > VirtioGetDevExtension(device, VirtioBLKDeviceConfig)->capacity * VIRTIO_BLK_SECTOR_SIZE)
            return NX_EINVAL;
        if (VirtioBLKReadWrite(
                device,
                off / VIRTIO_BLK_SECTOR_SIZE,
                buf,
                len / VIRTIO_BLK_SECTOR_SIZE, NX_False
        ) != NX_EOK)
            return NX_ERROR;
        if (outLen) *outLen = len;
        return NX_EOK;
    }
    else
    {
        // Not aligned, read into internal buffer first
        NX_Size sector          = off / VIRTIO_BLK_SECTOR_SIZE;
        NX_Size sectorOffset    = off % VIRTIO_BLK_SECTOR_SIZE;
        NX_Size elen            = len + sectorOffset;
        NX_Size bufSize         = NX_ALIGN_UP(elen, VIRTIO_BLK_SECTOR_SIZE);
        char    *internalBuffer = NX_NULL;
        if (sector + bufSize / VIRTIO_BLK_SECTOR_SIZE > VirtioGetDevExtension(device, VirtioBLKDeviceConfig)->capacity)
            return NX_EINVAL;
        if (bufSize > NX_PAGE_SIZE)
            internalBuffer = NX_PageAlloc(NX_PAGE_ALIGNUP(bufSize) / NX_PAGE_SIZE);
        else
            internalBuffer = NX_MemAlloc(bufSize);
        if (!internalBuffer) return NX_ENOMEM;
        if (VirtioBLKReadWrite(
                device,
                sector,
                internalBuffer,
                bufSize / VIRTIO_BLK_SECTOR_SIZE, NX_False
        ) != NX_EOK)
        {
            NX_MemFree(internalBuffer);
            return NX_ERROR;
        }
        NX_MemCopy(buf, internalBuffer + sectorOffset, len);
        if (outLen) *outLen = len;
        if (bufSize > NX_PAGE_SIZE)
            NX_PageFree(internalBuffer);
        else
            NX_MemFree(internalBuffer);
        return NX_EOK;
    }
}

NX_PRIVATE NX_Error write(struct NX_Device *device, void *buf, NX_Offset off, NX_Size len, NX_Size *outLen)
{
    if (VirtioGetExtension(device)->flags & VIRTIO_BLK_FLAGS_READ_ONLY)
        return NX_EINVAL;

    if (buf == NX_NULL) return NX_EINVAL;
    if (len == 0) return NX_EINVAL;
    buf = VirtioVirt2Phy(buf);
    if (off % VIRTIO_BLK_SECTOR_SIZE == 0 && len % VIRTIO_BLK_SECTOR_SIZE == 0)
    {
        if (off + len > VirtioGetDevExtension(device, VirtioBLKDeviceConfig)->capacity * VIRTIO_BLK_SECTOR_SIZE)
            return NX_EINVAL;
        // Perfect aligned, directly pass buffer to device without internal dirty buf
        if (VirtioBLKReadWrite(
                device,
                off / VIRTIO_BLK_SECTOR_SIZE,
                buf,
                len / VIRTIO_BLK_SECTOR_SIZE, NX_True
        ) != NX_EOK)
            return NX_ERROR;
        if (outLen) *outLen = len;
        return NX_EOK;
    }
    else
    {
        // Not aligned, read and write
        NX_Size sector          = off / VIRTIO_BLK_SECTOR_SIZE;
        NX_Size sectorOffset    = off % VIRTIO_BLK_SECTOR_SIZE;
        NX_Size elen            = len + sectorOffset;
        NX_Size bufSize         = NX_ALIGN_UP(elen, VIRTIO_BLK_SECTOR_SIZE);
        char    *internalBuffer = NX_NULL;
        if (bufSize > NX_PAGE_SIZE)
            internalBuffer = NX_PageAlloc(NX_PAGE_ALIGNUP(bufSize) / NX_PAGE_SIZE);
        else
            internalBuffer = NX_MemAlloc(bufSize);
        if (!internalBuffer) return NX_ENOMEM;
        // Read
        if (VirtioBLKReadWrite(
                device,
                sector,
                internalBuffer,
                bufSize / VIRTIO_BLK_SECTOR_SIZE, NX_False
        ) != NX_EOK)
        {
            NX_MemFree(internalBuffer);
            return NX_ERROR;
        }
        // Write to read buffer
        NX_MemCopy(internalBuffer + sectorOffset, buf, len);
        // Write to disk
        if (VirtioBLKReadWrite(
                device,
                sector,
                internalBuffer,
                bufSize / VIRTIO_BLK_SECTOR_SIZE, NX_True
        ) != NX_EOK)
        {
            NX_MemFree(internalBuffer);
            return NX_ERROR;
        }
        if (outLen) *outLen = len;
        if (bufSize > NX_PAGE_SIZE)
            NX_PageFree(internalBuffer);
        else
            NX_MemFree(internalBuffer);
        return NX_EOK;
    }
}

NX_PRIVATE NX_Error control(struct NX_Device *device, NX_U32 cmd, void *arg)
{
    NX_IoBlockInfo *info = NX_NULL;
    switch (cmd)
    {
    case NX_IO_BLOCK_INFO:
        info = (NX_IoBlockInfo *) arg;
        info->capacity   = VirtioGetDevExtension(device, VirtioBLKDeviceConfig)->capacity * VIRTIO_BLK_SECTOR_SIZE;
        info->blockSize  = VIRTIO_BLK_SECTOR_SIZE;
        info->blockCount = info->capacity / info->blockSize + 1;
        break;
    case NX_IO_BLOCK_SYNC:
        VirtioBLKFlush(device);
        break;
    default:
        break;
    }
    return NX_EOK;
}

// VirtioBLK driver encourages user using 512bytes aligned offset/size to improve the performance
NX_PRIVATE NX_DriverOps virtioBLKDriverOps = {
        .open = open,
        .close = close,
        .read = read,
        .write = write,
        .control = control
};

NX_PRIVATE NX_Error VirtioBLKInit(NX_Driver *driver, NX_Device *device)
{
    return NX_EOK;
}


NX_PRIVATE VirtioDriverTableEntity virtioBLKDriverTableEntity = {
        .deviceId = VIRTIO_DEVICE_ID_BLOCK,
        .driverName = "virtio-blk",
        .deviceNameTemplate = "virtio_blk_%d",
        .deviceType = NX_DEVICE_TYPE_BLOCK,
        .driverOps = &virtioBLKDriverOps,
        .initHandler = VirtioBLKInit,
        .cleanUpHandler = NX_NULL
};

VirtioDriverTableEntity *VirtioDriverTableEntityForVirtioBLK(void)
{
    return &virtioBLKDriverTableEntity;
}

#else
#include "../virtio.h"

VirtioDriverTableEntity *VirtioDriverTableEntityForVirtioBLK(void)
{
    return NX_NULL;
}

#endif // CONFIG_NX_DRIVER_VIRTIO_BLK
